import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 生成模拟数据
np.random.seed(0)
times = pd.date_range(start='2023-01-01 09:00:00', periods=240, freq='T')
prices = 100 + -5 * np.random.randn(240).cumsum()  # 价格在100附近随机波动
sid_price_df = pd.DataFrame({'datetime': times, 'close': prices})

# 计算5分钟和21分钟的均线（确保从有足够数据点的索引开始）
sid_price_df['5min_avg'] = sid_price_df['close'].rolling(window=5, min_periods=5).mean()
sid_price_df['21min_avg'] = sid_price_df['close'].rolling(window=21, min_periods=21).mean()

# 初始化买卖信号列并填充为0（从有足够数据点的索引开始）
sid_price_df['buy_signal'] = 0
sid_price_df['sell_signal'] = 0

# 找出买卖信号（从有足够数据点的索引开始）
start_idx = max(5 - 1, 21 - 1)  # 实际上这里取5，因为需要5个数据点来计算5分钟均线

for i in range(start_idx, len(sid_price_df)):
    # 检查前一个值是否为NaN（即确保有足够的数据点来计算均线）
    if not (pd.isnull(sid_price_df.at[i - 1, '5min_avg']) or pd.isnull(sid_price_df.at[i - 1, '21min_avg'])):
        # 卖出信号
        if sid_price_df.iloc[i - 1]['5min_avg'] > sid_price_df.iloc[i - 1]['21min_avg'] and sid_price_df.iloc[i - 1]['5min_avg'] < sid_price_df.iloc[i - 1]['21min_avg']:
            sid_price_df.iloc[i - 1]['sell_signal'] = 1
        
        # if sid_price_df.at[i - 1, '5min_avg'] > sid_price_df.at[i - 1, '21min_avg'] and sid_price_df.at[i, '5min_avg'] < sid_price_df.at[i, '21min_avg']:
        #     sid_price_df.at[i, 'sell_signal'] = 1
            # 买入信号
        if sid_price_df.at[i - 1, '5min_avg'] < sid_price_df.at[i - 1, '21min_avg'] and sid_price_df.at[i, '5min_avg'] > sid_price_df.at[i, '21min_avg']:
            sid_price_df.at[i, 'buy_signal'] = 1
            # 绘制图形
plt.figure(figsize=(12, 6))
plt.plot(sid_price_df['datetime'], sid_price_df['close'], label='close', alpha=0.7)
plt.plot(sid_price_df['datetime'], sid_price_df['5min_avg'], label='5-minute MA', linestyle='--', alpha=0.7)
plt.plot(sid_price_df['datetime'], sid_price_df['21min_avg'], label='21-minute MA', linestyle='-.', alpha=0.7)

# 标记买卖信号
buy_signals = sid_price_df[sid_price_df['buy_signal'] == 1]
sell_signals = sid_price_df[sid_price_df['sell_signal'] == 1]

plt.scatter(buy_signals['datetime'], buy_signals['close'], marker='^', color='g', label='Buy Signal')
plt.scatter(sell_signals['datetime'], sell_signals['close'], marker='v', color='r', label='Sell Signal')

# 设置图例和标题
plt.legend(loc='upper left')
plt.title('Dual Moving Average Strategy')
plt.xlabel('Datetime')
plt.ylabel('close')
plt.grid(True)
plt.xticks(rotation=45)

# 显示图形
plt.show()

x = 1